import os
import json
import random
import numpy as np
from collections import Counter, defaultdict
from PIL import Image, ImageDraw, ImageFont
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import font_manager

# ================== 目录结构与路径 ==================
target_dir = '/home/guoyi/Dataset/aier_processed/'
annotation_path = os.path.join(target_dir, 'annotations_v1.json')
experiments_dir = './experiments/cold_start'
os.makedirs(target_dir, exist_ok=True)
os.system(f"rm -rf {experiments_dir}/*")  # 清理旧实验目录
image_dir = os.path.join(target_dir, 'images')
visual_dir = os.path.join(experiments_dir, 'visual')
os.makedirs(visual_dir, exist_ok=True)
font_path = './SarasaMonoCL-Regular.ttf'  # 中文字体路径
report_path = os.path.join(experiments_dir, 'report.json')
report_log_path = os.path.join(experiments_dir, 'report_log.json')
analyse_image_dir = os.path.join(experiments_dir, 'analyse_images')
os.makedirs(analyse_image_dir, exist_ok=True)

ch_font_default = font_manager.FontProperties(fname=font_path, size=18)
matplotlib.rcParams['axes.unicode_minus'] = False  # 负号也能显示

def draw_bar_xtick_rightalign(
    counter_dict, 
    x_label, 
    y_label, 
    title, 
    save_path, 
    font_path, 
    max_bar_width=0.7, 
    x_tick_rotation=60,
    font_size=18,
    dpi=600
):
    """x轴标签右对齐且对齐bar中心，适合标签长短不一，防止重叠"""
    # 高到低排序
    sorted_items = sorted(counter_dict.items(), key=lambda x: x[1], reverse=True)
    keys = [k for k, _ in sorted_items]
    values = [v for _, v in sorted_items]
    num_bars = len(keys)
    fig_w = max(8, num_bars * 0.45)
    fig_h = 7

    ch_font = font_manager.FontProperties(fname=font_path, size=font_size)
    plt.figure(figsize=(fig_w, fig_h))
    x = np.arange(num_bars)
    bars = plt.bar(x, values, width=max_bar_width)
    plt.xlabel(x_label, fontproperties=ch_font)
    plt.ylabel(y_label, fontproperties=ch_font)
    plt.title(title, fontproperties=ch_font)
    # 手动右对齐所有xtick标签，且anchor对准bar中心
    plt.xticks(
        x, keys, rotation=x_tick_rotation, ha='right', va='bottom', fontproperties=ch_font, rotation_mode='anchor'
    )
    plt.yticks(fontproperties=ch_font)
    plt.tight_layout()
    plt.savefig(save_path, dpi=dpi)
    plt.close()

def draw_bar_xnatural_lesionnum(
    lesion_num_counter, 
    x_label, 
    y_label, 
    title, 
    save_path, 
    font_path, 
    font_size=24,
    dpi=600
):
    """病灶数直方图，x自然顺序，>10合并，标签大字体"""
    lesion_nums = sorted([k for k in lesion_num_counter if isinstance(k, int) or str(k).isdigit()])
    bar_counts = []
    xtick_labels = []
    more_than_10 = 0
    for n in range(1, 11):
        count = lesion_num_counter.get(n, 0)
        bar_counts.append(count)
        xtick_labels.append(str(n))
    for k in lesion_num_counter:
        if (isinstance(k, int) or str(k).isdigit()) and int(k) > 10:
            more_than_10 += lesion_num_counter[k]
        elif isinstance(k, str) and k.startswith('>'):
            more_than_10 += lesion_num_counter[k]
    bar_counts.append(more_than_10)
    xtick_labels.append(">10")

    num_bars = len(bar_counts)
    fig_w = max(10, num_bars * 0.7)
    fig_h = 6
    ch_font = font_manager.FontProperties(fname=font_path, size=font_size)
    plt.figure(figsize=(fig_w, fig_h))
    x = np.arange(num_bars)
    plt.bar(x, bar_counts, width=0.7)
    plt.xlabel(x_label, fontproperties=ch_font)
    plt.ylabel(y_label, fontproperties=ch_font)
    plt.title(title, fontproperties=ch_font)
    plt.xticks(x, xtick_labels, fontproperties=ch_font)
    plt.yticks(fontproperties=ch_font)
    plt.tight_layout()
    plt.savefig(save_path, dpi=dpi)
    plt.close()

# ================== 加载数据 ==================
with open(annotation_path, 'r', encoding='utf-8') as f:
    annotations = json.load(f)

# ================== 1. 样本图片可视化（不变） ==================
visual_num = 5  # 每类可视化图片数量
visualized_images = []
class_to_images = {}
for img_name, info in annotations.items():
    class_to_images.setdefault(info['image_label'], []).append(img_name)

for cls, img_list in class_to_images.items():
    selected = random.sample(img_list, min(visual_num, len(img_list)))
    for img_name in selected:
        img_path = os.path.join(image_dir, img_name)
        try:
            img = Image.open(img_path).convert("RGB")
            draw = ImageDraw.Draw(img)
            font = ImageFont.truetype(font_path, 16)
            for lesion in annotations[img_name]['lesions']:
                bbox = lesion['bbox']
                x0 = int(bbox[0] * 224)
                y0 = int(bbox[1] * 224)
                x1 = int(bbox[2] * 224)
                y1 = int(bbox[3] * 224)
                draw.rectangle([x0, y0, x1, y1], outline="red", width=2)
                draw.text((x0, y0), lesion['name'], font=font, fill="yellow")
            out_path = os.path.join(visual_dir, img_name)
            img.save(out_path)
            visualized_images.append({'image_name': img_name, 'image_path': out_path})
        except Exception as e:
            print(f"可视化失败: {img_name}, 错误: {e}")

# ================== 2. 各类别图片数量分布 (x右对齐) ==================
class_count = {cls: len(imgs) for cls, imgs in class_to_images.items()}
bar_path1 = os.path.join(analyse_image_dir, "class_count.png")
draw_bar_xtick_rightalign(
    class_count,
    x_label="类别",
    y_label="图片数量",
    title="各类别图片数量分布",
    save_path=bar_path1,
    font_path=font_path,
    x_tick_rotation=60,
    font_size=18
)

# ================== 3. 每张图片lesion数量分布 (自然序+大于10合并) ==================
lesion_nums = [len(info['lesions']) for info in annotations.values()]
lesion_num_counter_raw = dict(Counter(lesion_nums))
# 统计信息
lesion_num_stat = {
    '平均数': float(np.mean(lesion_nums)),
    '中位数': float(np.median(lesion_nums)),
    '标准差': float(np.std(lesion_nums)),
    '最小值': int(np.min(lesion_nums)),
    '最大值': int(np.max(lesion_nums)),
}
# 绘图
bar_path2 = os.path.join(analyse_image_dir, "lesion_num_hist.png")
draw_bar_xnatural_lesionnum(
    lesion_num_counter_raw,
    x_label="每张图片的病灶数",
    y_label="图片数",
    title="图片中病灶数量分布",
    save_path=bar_path2,
    font_path=font_path,
    font_size=24
)

# ================== 4. 病灶名称分布分析 (x右对齐) ==================
name_counter = Counter()
for info in annotations.values():
    for lesion in info['lesions']:
        name_counter[lesion['name']] += 1
name_counter_sorted = dict(sorted(name_counter.items(), key=lambda x: x[1], reverse=True))
bar_path3 = os.path.join(analyse_image_dir, "lesion_name_count.png")
draw_bar_xtick_rightalign(
    name_counter_sorted,
    x_label="病灶名称",
    y_label="数量",
    title="病灶名称分布",
    save_path=bar_path3,
    font_path=font_path,
    x_tick_rotation=60,
    font_size=18
)

# ================== 5. 框重叠分析 (不变，高清) ==================
def overlap_min_area(b1, b2):
    # 输入: [xmin, ymin, xmax, ymax], 归一化
    x1 = max(b1[0], b2[0])
    y1 = max(b1[1], b2[1])
    x2 = min(b1[2], b2[2])
    y2 = min(b1[3], b2[3])
    inter_w = max(0, x2 - x1)
    inter_h = max(0, y2 - y1)
    inter_area = inter_w * inter_h
    area1 = (b1[2] - b1[0]) * (b1[3] - b1[1])
    area2 = (b2[2] - b2[0]) * (b2[3] - b2[1])
    min_area = min(area1, area2)
    if min_area == 0:
        return 0
    return inter_area / min_area


max_overlap_vals = []
overlap_gt08_counter = 0
for img_name, info in annotations.items():
    lesions = info['lesions']
    max_overlap = 0
    name2bbox = defaultdict(list)
    for lesion in lesions:
        name2bbox[lesion['name']].append(lesion['bbox'])
    for name, bboxes in name2bbox.items():
        for i in range(len(bboxes)):
            for j in range(i + 1, len(bboxes)):
                iou = overlap_min_area(bboxes[i], bboxes[j])
                if iou > max_overlap:
                    max_overlap = iou
    max_overlap_vals.append(max_overlap)
    if max_overlap > 0.8:
        overlap_gt08_counter += 1

overlap_stat = {
    "平均重叠比": float(np.mean(max_overlap_vals)),
    "中位数": float(np.median(max_overlap_vals)),
    "标准差": float(np.std(max_overlap_vals)),
    "最大重叠比大于0.8的图片数": overlap_gt08_counter,
    "总图片数": len(max_overlap_vals)
}
bins = np.linspace(0, 1, 11)
hist, bin_edges = np.histogram(max_overlap_vals, bins=bins)
overlap_hist = {
    "区间": [f"{round(bin_edges[i],1)}-{round(bin_edges[i+1],1)}" for i in range(len(bin_edges)-1)],
    "图片数": hist.tolist()
}
# 重叠比分布图
fig_w = 10 if len(overlap_hist["区间"]) < 15 else len(overlap_hist["区间"]) * 0.7
ch_font_hist = font_manager.FontProperties(fname=font_path, size=18)
bar_path4 = os.path.join(analyse_image_dir, "max_overlap_hist.png")
plt.figure(figsize=(fig_w, 6))
plt.bar(range(len(overlap_hist['区间'])), overlap_hist['图片数'])
plt.xlabel("图片中最大重叠比", fontproperties=ch_font_hist)
plt.ylabel("图片数", fontproperties=ch_font_hist)
plt.title("图片最大重叠比分布", fontproperties=ch_font_hist)
plt.xticks(range(len(overlap_hist['区间'])), overlap_hist['区间'], rotation=45, fontproperties=ch_font_hist)
plt.yticks(fontproperties=ch_font_hist)
plt.tight_layout()
plt.savefig(bar_path4, dpi=600)
plt.close()
# ========== 7. 高重叠比分析（重叠比>0.6） ==========
high_overlap_images = []
high_overlap_labels = []
high_overlap_names = []

# 这里要重新统计图片中最大重叠比以及其name
for img_name, info in annotations.items():
    lesions = info['lesions']
    name2bbox = defaultdict(list)
    for lesion in lesions:
        name2bbox[lesion['name']].append(lesion['bbox'])
    # 找到该图片内最大重叠比及其对应name
    max_overlap = 0
    overlap_names = []
    for name, bboxes in name2bbox.items():
        for i in range(len(bboxes)):
            for j in range(i + 1, len(bboxes)):
                iou = overlap_min_area(bboxes[i], bboxes[j])
                if iou > max_overlap:
                    max_overlap = iou
                    overlap_names = [name]
                elif iou == max_overlap and iou > 0:
                    overlap_names.append(name)
    if max_overlap > 0.6:
        high_overlap_images.append(img_name)
        high_overlap_labels.append(info["image_label"])
        # 可能多个name达到最大重叠比，这里去重
        high_overlap_names += list(set(overlap_names))

# image_label分布
high_overlap_label_counter = dict(Counter(high_overlap_labels))
bar_path_label = os.path.join(analyse_image_dir, "high_overlap_image_label.png")
draw_bar_xtick_rightalign(
    high_overlap_label_counter,
    x_label="类别",
    y_label="高重叠图片数",
    title="高重叠比图片中各类别分布",
    save_path=bar_path_label,
    font_path=font_path,
    x_tick_rotation=60,
    font_size=18
)

# lesion name分布
high_overlap_name_counter = dict(Counter(high_overlap_names))
bar_path_name = os.path.join(analyse_image_dir, "high_overlap_lesion_name.png")
draw_bar_xtick_rightalign(
    high_overlap_name_counter,
    x_label="病灶名称",
    y_label="高重叠lesion数",
    title="高重叠比lesion名称分布",
    save_path=bar_path_name,
    font_path=font_path,
    x_tick_rotation=60,
    font_size=18
)

# ================== 7. 写入report与log ==================
report = {
    "class_count_chart": bar_path1,
    "lesion_num_stat": lesion_num_stat,
    "lesion_num_hist_chart": bar_path2,
    "lesion_name_counter": name_counter_sorted,
    "lesion_name_chart": bar_path3,
    "overlap_stat": overlap_stat,
    "overlap_hist_chart": bar_path4,
    "high_overlap_label_chart": bar_path_label,
    "high_overlap_name_chart": bar_path_name,
}

with open(report_path, 'w', encoding='utf-8') as f:
    json.dump(report, f, ensure_ascii=False, indent=2)

report_log = {
    "visual_image": visualized_images,
    "class_count": class_count,
    "lesion_num_counter": lesion_num_counter_raw,
    "max_overlap_per_image": max_overlap_vals,
    "lesion_name_counter": name_counter_sorted,
}
with open(report_log_path, 'w', encoding='utf-8') as f:
    json.dump(report_log, f, ensure_ascii=False, indent=2)

print(f"分析完成！报告见 {report_path}，详细log见 {report_log_path}")
